# -*- coding: utf-8 -*-
#
# Copyright (c) 2025 Huawei Technologies Co., Ltd.
# openUBMC is licensed under Mulan PSL v2.
# You can use this software according to the terms and conditions of the Mulan PSL v2.
# You may obtain a copy of Mulan PSL v2 at:
#         http://license.coscl.org.cn/MulanPSL2
# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
# See the Mulan PSL v2 for more details.

"""
发射器索引系统

提供高效的发射器查询和匹配机制，支持多种索引策略。
"""

from typing import List, Dict, Set, Optional, Type, Any
from dataclasses import dataclass, field
from abc import ABC, abstractmethod
from ...ir import IRNode
from ...ir.expressions import (
    IRCall,
    IRBinOp,
    IRUnaryOp,
    IRCompare,
    IRBoolOp,
    IRName,
    IRAttribute,
)
from .emitter import BaseEmitter
from .context import CodeGenContext


@dataclass
class EmitterMatchResult:
    """发射器匹配结果

    包含匹配的发射器和相关的缓存数据。
    """

    emitter: BaseEmitter
    priority: int
    cached_data: Dict[str, Any] = field(default_factory=dict)
    can_continue: bool = True  # 是否可以继续匹配其他发射器


class EmitterMatcher(ABC):
    """发射器匹配器基类

    定义统一的匹配接口，支持不同的匹配策略。
    """

    @abstractmethod
    def can_match(self, node: IRNode, context: CodeGenContext) -> bool:
        """检查是否可以匹配该节点"""
        pass

    @abstractmethod
    def match(self, node: IRNode, context: CodeGenContext) -> List[EmitterMatchResult]:
        """匹配发射器并返回结果"""
        pass

    @abstractmethod
    def add_emitter(self, emitter: BaseEmitter) -> None:
        """添加发射器到匹配器"""
        pass

    @abstractmethod
    def remove_emitter(self, emitter: BaseEmitter) -> None:
        """从匹配器中移除发射器"""
        pass


class NodeTypeIndex(EmitterMatcher):
    """基于节点类型的索引"""

    def __init__(self):
        self._type_to_emitters: Dict[Type[IRNode], List[BaseEmitter]] = {}
        self._all_emitters: Set[BaseEmitter] = set()

    def can_match(self, node: IRNode, context: CodeGenContext) -> bool:
        """检查是否有匹配的类型索引"""
        node_type = type(node)
        return node_type in self._type_to_emitters

    def match(self, node: IRNode, context: CodeGenContext) -> List[EmitterMatchResult]:
        """基于节点类型匹配发射器"""
        node_type = type(node)
        emitters = self._type_to_emitters.get(node_type, [])

        results = []
        for emitter in emitters:
            if emitter.can_emit(node, context):
                results.append(
                    EmitterMatchResult(
                        emitter=emitter,
                        priority=emitter.get_priority(),
                        cached_data={"match_type": "node_type", "node_type": node_type},
                    )
                )

        return results

    def add_emitter(self, emitter: BaseEmitter) -> None:
        """添加发射器到类型索引"""
        if emitter in self._all_emitters:
            return

        self._all_emitters.add(emitter)

        # 获取发射器支持的节点类型
        supported_types = self._get_supported_types(emitter)
        for node_type in supported_types:
            if node_type not in self._type_to_emitters:
                self._type_to_emitters[node_type] = []
            self._type_to_emitters[node_type].append(emitter)

    def remove_emitter(self, emitter: BaseEmitter) -> None:
        """从类型索引中移除发射器"""
        if emitter not in self._all_emitters:
            return

        self._all_emitters.remove(emitter)

        # 从所有类型索引中移除
        for emitters_list in self._type_to_emitters.values():
            if emitter in emitters_list:
                emitters_list.remove(emitter)

    def _get_supported_types(self, emitter: BaseEmitter) -> List[Type[IRNode]]:
        """获取发射器支持的节点类型"""
        if hasattr(emitter, "supported_types"):
            return getattr(emitter, "supported_types")

        return []


class FunctionCallIndex(EmitterMatcher):
    """基于函数调用的索引"""

    def __init__(self):
        self._function_to_emitters: Dict[str, List[BaseEmitter]] = {}
        self._all_emitters: Set[BaseEmitter] = set()

    def can_match(self, node: IRNode, context: CodeGenContext) -> bool:
        """检查是否为函数调用节点"""
        return isinstance(node, IRCall)

    def match(self, node: IRNode, context: CodeGenContext) -> List[EmitterMatchResult]:
        """基于函数名匹配发射器"""
        if not isinstance(node, IRCall):
            return []

        func_name = self._get_function_name(node)
        if not func_name:
            return []

        emitters = self._function_to_emitters.get(func_name, [])
        results = []

        for emitter in emitters:
            if emitter.can_emit(node, context):
                # 缓存函数名和参数类型信息
                cached_data = {
                    "match_type": "function_call",
                    "function_name": func_name,
                    "symbol": None,
                }

                # 如果是模板函数，缓存符号查找结果
                if hasattr(emitter, "_lookup_symbol"):
                    _lookup_symbol = getattr(emitter, "_lookup_symbol")
                    args_types = self._get_args_types(node, context)
                    symbol = _lookup_symbol(node, func_name, context, args_types)
                    cached_data.update({"symbol": symbol, "args_types": args_types})

                results.append(
                    EmitterMatchResult(
                        emitter=emitter,
                        priority=emitter.get_priority(),
                        cached_data=cached_data,
                    )
                )

        return results

    def add_emitter(self, emitter: BaseEmitter) -> None:
        """添加发射器到函数索引"""
        if emitter in self._all_emitters:
            return

        self._all_emitters.add(emitter)

        # 获取发射器支持的函数名
        supported_functions = self._get_supported_functions(emitter)
        for func_name in supported_functions:
            if func_name not in self._function_to_emitters:
                self._function_to_emitters[func_name] = []
            self._function_to_emitters[func_name].append(emitter)

    def remove_emitter(self, emitter: BaseEmitter) -> None:
        """从函数索引中移除发射器"""
        if emitter not in self._all_emitters:
            return

        self._all_emitters.remove(emitter)

        # 从所有函数索引中移除
        for emitters_list in self._function_to_emitters.values():
            if emitter in emitters_list:
                emitters_list.remove(emitter)

    def _get_function_name(self, node: IRCall) -> Optional[str]:
        """获取函数调用的名称"""
        if isinstance(node.func, IRName):
            return node.func.id
        elif isinstance(node.func, IRAttribute):
            return node.func.attr
        return None

    def _get_supported_functions(self, emitter: BaseEmitter) -> List[str]:
        """获取发射器支持的函数名"""
        if hasattr(emitter, "supported_functions"):
            return getattr(emitter, "supported_functions")
        return []

    def _get_args_types(self, node: IRCall, context: CodeGenContext) -> List[Any]:
        """获取函数调用参数的类型信息"""
        args_types = []
        for arg in node.args or []:
            if hasattr(arg, "get_extension_param"):
                type_info = arg.get_extension_param("inferred_type")
                args_types.append(type_info)
            else:
                args_types.append(None)
        return args_types


class OperatorIndex(EmitterMatcher):
    """基于操作符的索引"""

    def __init__(self):
        self._operator_to_emitters: Dict[str, List[BaseEmitter]] = {}
        self._all_emitters: Set[BaseEmitter] = set()

    def can_match(self, node: IRNode, context: CodeGenContext) -> bool:
        """检查是否为操作符节点"""
        from ...ir.expressions.control import IRConditional

        return isinstance(
            node, (IRBinOp, IRUnaryOp, IRBoolOp, IRCompare, IRConditional)
        )

    def match(self, node: IRNode, context: CodeGenContext) -> List[EmitterMatchResult]:
        """基于操作符匹配发射器"""
        operators = self._get_node_operators(node)
        if not operators:
            return []

        results = []
        for op in operators:
            emitters = self._operator_to_emitters.get(op, [])
            for emitter in emitters:
                if emitter.can_emit(node, context):
                    cached_data = {
                        "match_type": "operator",
                        "operator": op,
                        "node_type": type(node).__name__,
                    }

                    # 对于操作符重载，缓存类型信息
                    if hasattr(emitter, "_get_node_type"):
                        if isinstance(node, (IRBinOp, IRCompare)):
                            _get_node_type = getattr(emitter, "_get_node_type")
                            left_type = _get_node_type(
                                getattr(node, "left", None), context
                            )
                            cached_data["left_type"] = left_type

                    results.append(
                        EmitterMatchResult(
                            emitter=emitter,
                            priority=emitter.get_priority(),
                            cached_data=cached_data,
                        )
                    )

        return results

    def add_emitter(self, emitter: BaseEmitter) -> None:
        """添加发射器到操作符索引"""
        if emitter in self._all_emitters:
            return

        self._all_emitters.add(emitter)

        # 获取发射器支持的操作符
        supported_operators = self._get_supported_operators(emitter)
        for op in supported_operators:
            if op not in self._operator_to_emitters:
                self._operator_to_emitters[op] = []
            self._operator_to_emitters[op].append(emitter)

    def remove_emitter(self, emitter: BaseEmitter) -> None:
        """从操作符索引中移除发射器"""
        if emitter not in self._all_emitters:
            return

        self._all_emitters.remove(emitter)

        # 从所有操作符索引中移除
        for emitters_list in self._operator_to_emitters.values():
            if emitter in emitters_list:
                emitters_list.remove(emitter)

    def _get_node_operators(self, node: IRNode) -> List[str]:
        """获取节点的操作符"""
        from ...ir.expressions.control import IRConditional

        if isinstance(node, (IRBinOp, IRUnaryOp, IRBoolOp)):
            return [node.op]
        elif isinstance(node, IRCompare):
            return node.ops
        elif isinstance(node, IRConditional):
            return ["?:"]  # 三元运算符的特殊标识
        return []

    def _get_supported_operators(self, emitter: BaseEmitter) -> List[str]:
        """获取发射器支持的操作符"""
        if hasattr(emitter, "supported_operators"):
            return getattr(emitter, "supported_operators")
        return []


class FallbackMatcher(EmitterMatcher):
    """回退匹配器

    用于处理无法通过索引快速匹配的发射器。
    """

    def __init__(self):
        self._emitters: List[BaseEmitter] = []

    def can_match(self, node: IRNode, context: CodeGenContext) -> bool:
        """总是可以匹配（作为回退）"""
        return True

    def match(self, node: IRNode, context: CodeGenContext) -> List[EmitterMatchResult]:
        """逐个检查所有发射器"""
        results = []
        for emitter in self._emitters:
            if emitter.can_emit(node, context):
                results.append(
                    EmitterMatchResult(
                        emitter=emitter,
                        priority=emitter.get_priority(),
                        cached_data={"match_type": "fallback"},
                    )
                )

        return results

    def add_emitter(self, emitter: BaseEmitter) -> None:
        """添加发射器到回退列表"""
        if emitter not in self._emitters:
            self._emitters.append(emitter)

    def remove_emitter(self, emitter: BaseEmitter) -> None:
        """从回退列表中移除发射器"""
        if emitter in self._emitters:
            self._emitters.remove(emitter)


class EmitterIndexSystem:
    """发射器索引系统

    管理多种索引策略，提供高效的发射器查询。
    """

    def __init__(self):
        self._matchers: List[EmitterMatcher] = [
            NodeTypeIndex(),
            FunctionCallIndex(),
            OperatorIndex(),
            FallbackMatcher(),  # 回退匹配器放在最后
        ]
        self._all_emitters: Set[BaseEmitter] = set()

    def add_emitter(self, emitter: BaseEmitter) -> None:
        """添加发射器到索引系统"""
        if emitter in self._all_emitters:
            return

        self._all_emitters.add(emitter)

        # 将发射器添加到合适的匹配器中
        added_to_specific_matcher = False

        for matcher in self._matchers[:-1]:  # 排除回退匹配器
            try:
                matcher.add_emitter(emitter)
                if self._emitter_fits_matcher(emitter, matcher):
                    added_to_specific_matcher = True
            except Exception:
                # 如果添加失败，继续尝试其他匹配器
                continue

        # 如果没有添加到特定匹配器，添加到回退匹配器
        if not added_to_specific_matcher:
            self._matchers[-1].add_emitter(emitter)

    def remove_emitter(self, emitter: BaseEmitter) -> None:
        """从索引系统中移除发射器"""
        if emitter not in self._all_emitters:
            return

        self._all_emitters.remove(emitter)

        # 从所有匹配器中移除
        for matcher in self._matchers:
            matcher.remove_emitter(emitter)

    def find_emitters(
        self, node: IRNode, context: CodeGenContext
    ) -> List[EmitterMatchResult]:
        """查找匹配的发射器"""
        all_results = []

        # 按顺序尝试各个匹配器
        for matcher in self._matchers:
            if matcher.can_match(node, context):
                results = matcher.match(node, context)
                all_results.extend(results)

        # 去重并按优先级排序
        unique_results = self._deduplicate_results(all_results)
        unique_results.sort(key=lambda r: r.priority, reverse=True)

        return unique_results

    def find_best_emitter(
        self, node: IRNode, context: CodeGenContext
    ) -> Optional[EmitterMatchResult]:
        """查找最佳匹配的发射器"""
        results = self.find_emitters(node, context)
        return results[0] if results else None

    def get_all_emitters(self) -> List[BaseEmitter]:
        """获取所有注册的发射器"""
        return list(self._all_emitters)

    def clear(self) -> None:
        """清空所有发射器"""
        self._all_emitters.clear()
        for matcher in self._matchers:
            if hasattr(matcher, "_emitters"):
                getattr(matcher, "_emitters").clear()
            if hasattr(matcher, "_all_emitters"):
                getattr(matcher, "_all_emitters").clear()

    def _emitter_fits_matcher(
        self, emitter: BaseEmitter, matcher: EmitterMatcher
    ) -> bool:
        """检查发射器是否适合特定匹配器"""
        if isinstance(matcher, NodeTypeIndex):
            return hasattr(emitter, "supported_types") and getattr(
                emitter, "supported_types"
            )
        elif isinstance(matcher, FunctionCallIndex):
            return hasattr(emitter, "supported_functions") and getattr(
                emitter, "supported_functions"
            )
        elif isinstance(matcher, OperatorIndex):
            return hasattr(emitter, "supported_operators") and getattr(
                emitter, "supported_operators"
            )
        return False

    def _deduplicate_results(
        self, results: List[EmitterMatchResult]
    ) -> List[EmitterMatchResult]:
        """去除重复的匹配结果"""
        seen_emitters = set()
        unique_results = []

        for result in results:
            if result.emitter not in seen_emitters:
                seen_emitters.add(result.emitter)
                unique_results.append(result)

        return unique_results
